"""
预测模型模块
负责机器学习模型的训练和股价预测
"""

import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from datetime import datetime, timedelta
import streamlit as st


class StockPredictor:
    """股票价格预测器"""
    
    def __init__(self, sequence_length=10):
        """
        初始化预测器
        
        Args:
            sequence_length: 时序数据长度，默认10天
        """
        self.sequence_length = sequence_length
        self.model = None
        self.scaler_X = None
        self.scaler_y = None
        self.feature_columns = ['open', 'high', 'low', 'volume', 'ma5', 'ma20', 'rsi', 'previous_close']
        
    def prepare_data(self, df):
        """
        准备训练数据
        
        Args:
            df: 包含技术指标的股票数据
            
        Returns:
            tuple: (X_scaled, y_scaled, feature_data, target_data)
        """
        # 选择特征
        X = df[self.feature_columns].values
        y = df['close'].values
        
        # 数据归一化
        self.scaler_X = MinMaxScaler()
        self.scaler_y = MinMaxScaler()
        
        X_scaled = self.scaler_X.fit_transform(X)
        y_scaled = self.scaler_y.fit_transform(y.reshape(-1, 1)).flatten()
        
        return X_scaled, y_scaled, X, y
    
    def create_sequences(self, X, y):
        """
        创建时序数据序列
        
        Args:
            X: 特征数据
            y: 目标数据
            
        Returns:
            tuple: (X_sequences, y_sequences)
        """
        X_seq, y_seq = [], []
        for i in range(self.sequence_length, len(X)):
            X_seq.append(X[i-self.sequence_length:i])
            y_seq.append(y[i])
        return np.array(X_seq), np.array(y_seq)
    
    def train(self, df, test_ratio=0.2):
        """
        训练模型
        
        Args:
            df: 股票数据
            test_ratio: 测试集比例
            
        Returns:
            dict: 训练结果和评估指标
        """
        # 准备数据
        X_scaled, y_scaled, X_raw, y_raw = self.prepare_data(df)
        
        # 创建序列
        X_seq, y_seq = self.create_sequences(X_scaled, y_scaled)
        
        if len(X_seq) < 20:
            raise ValueError("数据量不足，无法训练模型")
        
        # 分割训练测试集
        train_size = int(len(X_seq) * (1 - test_ratio))
        
        X_train = X_seq[:train_size]
        X_test = X_seq[train_size:]
        y_train = y_seq[:train_size]
        y_test = y_seq[train_size:]
        
        # 重塑数据用于线性回归
        X_train_flat = X_train.reshape(X_train.shape[0], -1)
        X_test_flat = X_test.reshape(X_test.shape[0], -1)
        
        # 训练模型
        self.model = LinearRegression()
        self.model.fit(X_train_flat, y_train)
        
        # 预测
        y_pred_train = self.model.predict(X_train_flat)
        y_pred_test = self.model.predict(X_test_flat)
        
        # 反归一化
        y_train_real = self.scaler_y.inverse_transform(y_train.reshape(-1, 1)).flatten()
        y_test_real = self.scaler_y.inverse_transform(y_test.reshape(-1, 1)).flatten()
        y_pred_train_real = self.scaler_y.inverse_transform(y_pred_train.reshape(-1, 1)).flatten()
        y_pred_test_real = self.scaler_y.inverse_transform(y_pred_test.reshape(-1, 1)).flatten()
        
        # 计算评估指标
        train_metrics = {
            'mse': mean_squared_error(y_train_real, y_pred_train_real),
            'mae': mean_absolute_error(y_train_real, y_pred_train_real),
            'r2': r2_score(y_train_real, y_pred_train_real)
        }
        
        test_metrics = {
            'mse': mean_squared_error(y_test_real, y_pred_test_real),
            'mae': mean_absolute_error(y_test_real, y_pred_test_real),
            'r2': r2_score(y_test_real, y_pred_test_real)
        }
        
        # 获取测试期间的日期
        test_start_idx = len(df) - len(y_test_real)
        test_dates = df.index[test_start_idx:]
        
        results = {
            'train_metrics': train_metrics,
            'test_metrics': test_metrics,
            'y_test_real': y_test_real,
            'y_pred_test_real': y_pred_test_real,
            'test_dates': test_dates,
            'model_trained': True
        }
        
        return results
    
    def predict_future(self, df, days=20):
        """
        预测未来股价
        
        Args:
            df: 股票数据
            days: 预测天数
            
        Returns:
            dict: 预测结果
        """
        if self.model is None:
            raise ValueError("模型未训练，请先调用train()方法")
        
        # 准备最新数据
        X_scaled, _, _, _ = self.prepare_data(df)
        last_sequence = X_scaled[-self.sequence_length:]
        
        future_predictions = []
        current_sequence = last_sequence.copy()
        
        for i in range(days):
            # 预测下一天
            next_pred = self.model.predict(current_sequence.reshape(1, -1))[0]
            future_predictions.append(next_pred)
            
            # 更新序列
            new_row = current_sequence[-1].copy()
            new_row[0] = next_pred  # 开盘价 ≈ 前收盘价
            new_row[1] = next_pred * 1.02  # 最高价
            new_row[2] = next_pred * 0.98  # 最低价
            new_row[-1] = next_pred  # previous_close
            
            # 滚动更新序列
            current_sequence = np.vstack([current_sequence[1:], new_row])
        
        # 反归一化预测结果
        future_predictions_real = self.scaler_y.inverse_transform(
            np.array(future_predictions).reshape(-1, 1)
        ).flatten()
        
        # 生成未来日期
        last_date = df.index[-1]
        future_dates = []
        current_date = last_date
        
        for i in range(days):
            current_date += timedelta(days=1)
            # 跳过周末
            while current_date.weekday() >= 5:
                current_date += timedelta(days=1)
            future_dates.append(current_date)
        
        # 当前价格
        current_price = df['close'].iloc[-1]
        
        # 计算变化
        changes = future_predictions_real - current_price
        changes_pct = (changes / current_price) * 100
        
        # 预测结果
        prediction_df = pd.DataFrame({
            'date': future_dates,
            'predicted_price': future_predictions_real,
            'change_amount': changes,
            'change_percent': changes_pct
        })
        
        # 统计信息
        stats = {
            'current_price': current_price,
            'min_predicted': future_predictions_real.min(),
            'max_predicted': future_predictions_real.max(),
            'avg_predicted': future_predictions_real.mean(),
            'total_change': future_predictions_real[-1] - current_price,
            'total_change_pct': ((future_predictions_real[-1] - current_price) / current_price) * 100
        }
        
        # 投资建议
        if stats['avg_predicted'] > current_price * 1.05:
            suggestion = "看涨 📈"
        elif stats['avg_predicted'] < current_price * 0.95:
            suggestion = "看跌 📉"
        else:
            suggestion = "震荡 ➡️"
        
        return {
            'predictions': prediction_df,
            'stats': stats,
            'suggestion': suggestion,
            'future_dates': future_dates,
            'future_prices': future_predictions_real
        }
    
    def get_model_info(self):
        """获取模型信息"""
        if self.model is None:
            return {"status": "未训练"}
        
        return {
            "status": "已训练",
            "model_type": "线性回归",
            "sequence_length": self.sequence_length,
            "features": self.feature_columns,
            "feature_count": len(self.feature_columns)
        }